# -*- coding: utf-8 -*-
# @Time    : 2019/2/27 15:53
# @Author  : cj
# @File    : data_cleaner.py
import os
import os.path
import xml.etree.cElementTree as ET
from xml.etree import cElementTree
from utils.basic import cv_imread, cv_imwrite


def _check_rule_satisify(label: str, rule: list) -> bool:
    if len(rule) == 0:
        return True
    for prefix in rule:
        index = label.find(prefix)
        if index != 0:
            continue
        index = index + len(prefix)
        if index >= len(label) or label[index:].isdigit():
            return True
    return False


class DataItemCleaner:
    @staticmethod
    def del_label_with_blank(xml_path: str, label_filtering_rule: list) -> (list, int):
        label_set = set()
        obj_cnt = 0
        if xml_path.split('.')[-1] == 'xml':
            root_element = cElementTree.parse(xml_path)
            root = root_element.getroot()
            ob_list = root.findall('object')
            for ob in ob_list:
                name = ob.find('./name')
                name.text = name.text.replace(" ", "")
                if len(name.text) == 0:
                    continue
                if _check_rule_satisify(name.text, label_filtering_rule):
                    obj_cnt += 1
                    label_set.add(name.text)

            root_element.write(xml_path)
        return list(label_set), obj_cnt


class DatasetCleaner:
    """
    DatasetCleaner：用于清理一些图片和xml常见的问题,包括：1.去XML中object下name中带空格的字符串
    2.缺少xml的图片和缺少图片的XML
    3.用于检测图片filename和xlm里的filename不一致的情况
    4.xml的大小与图片大小不一致.min>=xmax,ymin>=ymax,xmin<0,xmax>=width,ymin<0,ymax>=height的情况
    5.用于检测xml中没有boundingbox的情况

    分别对应五个函数，当调用时会直接处理这五种情况
    """

    def __init__(self, path: str, label_filtering_rule: list):
        self._path = path
        self._rule = label_filtering_rule

    def del_label_with_blank(self):
        """
        去XML中object下name中带空格的字符串，程序输出去掉的个数
        """
        all_label_set = set()
        for file_idx in os.listdir(self._path):
            if file_idx.split('.')[-1] == 'xml':
                xml_path = os.path.join(self._path, file_idx)
                label_list, obj_cnt = DataItemCleaner.del_label_with_blank(xml_path, self._rule)
                if obj_cnt == 0:
                    os.remove(xml_path)
                    continue
                all_label_set = all_label_set | set(label_list)
        return list(all_label_set)

    def process_xml_without_picture(self):
        """
        find_noxmlpicture:找出缺少图片的XML,并删除
        path:包含图片的相应的xml的路径
        输出：no_xml，缺少xml的图片
              no_picture：缺少图片的XML
        """
        no_picture = []
        picture_path = []
        picture_full_path = []
        xml_path = []
        xml_full_path = []
        dataset_items = []
        for file_name in os.listdir(self._path):
            if file_name.split('.')[-1] == "jpg" or file_name.split('.')[-1] == "png" or \
                    file_name.split('.')[-1] == "jpeg":
                picture_path.append(file_name.split('.')[0])
                picture_full_path.append(file_name)
            elif file_name.split('.')[-1] == "xml":
                xml_path.append(file_name.split('.')[0])
                xml_full_path.append(file_name)

        for idx, xml_name in enumerate(xml_path):
            if xml_name not in picture_path:
                no_picture.append(xml_name)
            else:
                pic_idx = picture_path.index(xml_name)
                dataset_items.append({"image": os.path.join(self._path, picture_full_path[pic_idx]),
                                      "xml": os.path.join(self._path, xml_full_path[idx])})
        annotated_num = len(dataset_items)  # 所有已标注的图片数量
        for idx, pic_name in enumerate(picture_path):  # 得到所有没有标注的图片路径
            if pic_name not in xml_path:
                dataset_items.append({"image": os.path.join(self._path, picture_full_path[idx])})

        for delete_xml in no_picture:  # 删除没有图片的xml
            delete_xml_name = os.path.join(self._path, delete_xml + '.xml')
            os.remove(delete_xml_name)

        return annotated_num, len(picture_path), dataset_items

    def process_img_size_problem(self):
        """
        img_size_problem:处理以下情况：1.xml的大小与图片大小不一致，将xml的大小更正成图片的大小
        2.检测xmin>=xmax,ymin>=ymax,xmin<0,xmax>=width,ymin<0,ymax>=height的情况，如果出现直接删除boundingbox
        """
        for file_name in os.listdir(self._path):
            if file_name.split('.')[-1] == "jpg" or file_name.split('.')[-1] == "png" or \
                    file_name.split('.')[-1] == "jpeg":
                img_path = os.path.join(self._path, file_name)
                img = cv_imread(img_path)
                img_width = img.shape[1]
                img_height = img.shape[0]
                cv_imwrite(img_path, img)
                try:
                    xml_path = os.path.join(self._path, file_name.split('.')[0:-1][0] + ".xml")
                    tree = ET.parse(xml_path)
                    root = tree.getroot()
                    size = root.find('size')
                    xml_width = int(size.find('width').text)
                    xml_height = int(size.find('height').text)
                    # 更改xml中shape和图片shape不一致问题
                    if img_width != xml_width:
                        size.find('width').text = str(img_width)
                    if img_height != xml_height:
                        size.find('height').text = str(img_height)
                        # 检测xmin>=xmax,ymin>=ymax,xmin<0,xmax>=width,ymin<0,ymax>=height

                    for obj in root.iter('object'):
                        xml_box = obj.find('bndbox')
                        _xmin = int(xml_box.find('xmin').text)
                        _xmax = int(xml_box.find('xmax').text)
                        _ymin = int(xml_box.find('ymin').text)
                        _ymax = int(xml_box.find('ymax').text)
                        if _xmin >= _xmax:
                            root.remove(obj)
                            continue
                        if _ymin >= _ymax:
                            root.remove(obj)
                            continue
                        if _xmin <= 0:
                            root.remove(obj)
                            continue
                        if _xmax > img_width:
                            root.remove(obj)
                            continue
                        if _ymin <= 0:
                            root.remove(obj)
                            continue
                        if _ymax > img_height:
                            root.remove(obj)
                            continue

                    tree.write(xml_path)
                except FileNotFoundError:
                    continue

    def correlation_filename(self):
        """
        correlation_filename:当图片的filename和xlm里的filename不一致时，将xml中的filename
        改为图片的filename
        """
        for file_name in os.listdir(self._path):
            if file_name.split('.')[-1] == "jpg" or file_name.split('.')[-1] == "png" or \
                    file_name.split('.')[-1] == "jpeg":
                img_file_name = file_name
                xml_path = os.path.join(self._path, file_name.split('.')[0:-1][0] + ".xml")
                try:
                    tree = ET.parse(xml_path)
                    root = tree.getroot()
                    xml_file_name = root.find('./filename')
                    if str(xml_file_name.text) != img_file_name:
                        xml_file_name.text = img_file_name

                    tree.write(xml_path)
                except FileNotFoundError:
                    continue
